#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/6/18 17:12
# @Author  : CHEN Wang
# @Site    : 
# @File    : conn_wrappers.py
# @Software: PyCharm

import time
from functools import wraps
from quant_researcher.quant.project_tool.db_operator.my_mysql import get_mysql_conn
from quant_researcher.quant.project_tool.logger.my_logger import LOG


def conn_meddler(conn_getter):
    """
    自动给函数加入一个 conn 到关键字参数中，用完了 close
        如果关键字参数中不包含 conn 或找到的是 None，才会执行创建
        这样就不用手动管理 conn 的创建和关闭，which is 重复&蛮简单&精确&可复用
    :param function conn_getter: 如何生成这个 conn 的函数
    :return: 另一个 wrapper
    """
    def wrapper(func):
        @wraps(func)
        def inner(*args, **kwargs):
            close_needed = False
            if kwargs.get('conn') is None:
                LOG.info('未找到conn或者传入None，将建立，然后在用后close')
                close_needed = True
            try:
                conn = conn_getter()
                conn_got_ns = time.time()
                LOG.info(f'成功获得conn，获得时间={conn_got_ns}')
            except Exception as err:
                LOG.error(f'在尝试调用配置的conn获取函数时发生错误：{err}')
                raise err
            kwargs.update({'conn': conn})
            try:
                ans = func(*args, **kwargs)
            finally:
                if close_needed:
                    LOG.info(f'现在关闭连接(此连接获得时间={conn_got_ns})！')
                    conn.close()
            return ans
        return inner
    return wrapper


def set_conn_v2(func):
    """
    一个简单的，设置conn装饰器
        如果字典参数中没有conn，就自动加一个进去，先不管db是谁了，用默认的吧
    :param func func: 即将被装饰的函数
    :return: function，经过装饰的函数
    """
    @wraps(func)
    def wrapper(*args, **kwargs):
        conn = kwargs.get('conn')
        close_it = False
        if conn is None:
            conn = get_mysql_conn(
                host=kwargs.get('host'), port=kwargs.get('port'),
                user=kwargs.get('user'), password=kwargs.get('password'),
                database=kwargs.get('db_str')
            )
            close_it = True
        kwargs.update({'conn': conn})
        try:
            result = func(*args, **kwargs)
        finally:
            if close_it:
                conn.close()
        return result
    return wrapper